import os

from tqdm import tqdm

# os.environ['CUDA_VISIBLE_DEVICES'] = '0,1'

from dataset.dataset import UnpairedDataset,  set_random,SingleDataset
from utils.ForReaadConfig import parse_config
import matplotlib.pyplot as plt
from model import SIFA
from torch.utils.data import DataLoader
import torch
import numpy as np
import matplotlib
from metrics import assd_eval
from utils.myMetrics import Meterics

matplotlib.use('Agg')

def evaluate(epoch,config,sifa_model):
    '''
        脾（spleen）：4
        右肾（right kidney）：2
        左肾（left kidney）：3
        肝：1
    '''
    device = torch.device('cuda:{}'.format(config['train']['gpu']))
    test_path = config['test']['test_path']
    num_classes = config['test']['num_classes']
    sifa_model.eval()
    # test dataset
    test_dataset = SingleDataset(test_path)
    batch_size = config['test']['batch_size']
    test_loader = DataLoader(test_dataset, batch_size, shuffle=False)
    # test
    all_batch_assd = []
    myMeteric = Meterics(config['test']['num_classes'])
    with torch.no_grad():
        for it, (xt, xt_label) in tqdm(enumerate(test_loader)):
            xt = xt.to(device)
            xt_label = xt_label.numpy().squeeze().astype(np.uint8)

            output = sifa_model.test_seg(xt).detach()
            output = output.squeeze(0)
            output = torch.argmax(output, dim=0)
            output = output.cpu().numpy()
            output = output.squeeze()
            myMeteric.add_batch(xt_label, output)
            # print('{:.4f} th case dice MYO:{:.4f} LV:{:.4f} RV:{:.4f}'.format(it+1,one_case_dice[0],one_case_dice[1],one_case_dice[2]))
            # dicefile.write('file:{},{} th case dice:{}\n'.format(filename,it+1,one_case_dice))
            try:
                one_case_assd = assd_eval(output, xt_label, num_classes)
            except:
                continue
            all_batch_assd.append(one_case_assd)

    all_batch_assd = np.array(all_batch_assd)
    mean_assd = np.mean(all_batch_assd, axis=0)
    mean_dice = myMeteric.get_dice()
    mean_precision = myMeteric.get_precision()
    mean_recall = myMeteric.get_recall()
    # 输出
    std_assd = np.std(all_batch_assd, axis=0)
    print('------epoch:{}-----'.format(epoch))
    print('MYO||LV||RV')
    print('Dice mean:{}'.format(mean_dice))
    print('precision mean:{}'.format(mean_precision))
    print('recall mean:{}'.format(mean_recall))
    print('ASSD mean:{}'.format(mean_assd))
    print('ASSD std:{}'.format(std_assd))
    print('total mean assd:', np.mean(mean_assd))
    print('-----------')
# train
def train():
    # load config
    config = "/home/liukai/projects/SIFA/config/trainDebug.cfg"
    config = parse_config(config)
    # load data
    print(config)
    A_path = config['train']['a_path']
    B_path = config['train']['b_path']
    batch_size = config['train']['batch_size']

    trainset = UnpairedDataset(A_path, B_path)
    train_loader = DataLoader(trainset, batch_size,
                              shuffle=True, drop_last=True)
    # load exp_name
    exp_name = config['train']['exp_name']

    loss_cycle = []
    loss_seg = []
    # load model

    device = torch.device('cuda:{}'.format(config['train']['gpu']))
    print(device)

    # device = torch.device("cuda")
    # device = torch.device('cpu')
    sifa_model = SIFA(config).to(device)
    sifa_model.train()
    sifa_model.initialize()
    num_epochs = config['train']['num_epochs']
    # save_epoch = num_epochs // 10

    for epoch in tqdm(range(num_epochs)):
        sifa_model.train()
        for i, (A, A_label, B, _) in enumerate(train_loader):
            A = A.to(device).detach()
            B = B.to(device).detach()
            A_label = A_label.to(device).detach()
            sifa_model.update_GAN(A, B)
            sifa_model.update_seg(A, B, A_label)
            loss_cyclea, loss_cycleb, segloss = sifa_model.print_loss()
            loss_cycle.append(loss_cyclea + loss_cycleb)
            loss_seg.append(segloss)
            if i%40==0:
                print('---------------------')
                print("[{}/{}]".format(i,len(train_loader)))
                print('loss cycle A:', loss_cyclea)
                print('loss cycle B:', loss_cycleb)
                print('loss seg:', segloss)
        # ddfseg_model.update_lr() #no need for changing lr
        if (epoch + 1) % 10 == 0:
            evaluate(epoch,config,sifa_model)
            model_dir = config['train']['model_save_path']
            if (not os.path.exists(model_dir)):
                os.makedirs(model_dir)
            sifa_model.sample_image(epoch, config['train']['sample_save_path'],exp_name)
            torch.save(sifa_model.state_dict(),
                       '{}/model-{}.pth'.format(model_dir, epoch + 1))
        sifa_model.update_lr()

    print('train finished')
    loss_cycle = np.array(loss_cycle)
    loss_seg = np.array(loss_seg)
    np.savez('trainingloss.npz', loss_cycle, loss_seg)
    x = np.arange(0, loss_cycle.shape[0])
    plt.figure(1)
    plt.plot(x, loss_cycle, label='cycle loss of training')
    plt.legend()
    plt.xlabel('iterations')
    plt.ylabel('cycle loss')
    plt.savefig('cycleloss.jpg')
    plt.close()
    plt.figure(2)
    plt.plot(x, loss_seg, label='seg loss of training')
    plt.legend()
    plt.xlabel('iterations')
    plt.ylabel('seg loss')
    plt.savefig('segloss.jpg')
    plt.close()
    print('loss saved')


if __name__ == '__main__':
    set_random()
    train()
